#!/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plot
from math import sqrt
import heapq

"""
  kd树（k-dimension tree，k是指特征向量的维数），\
  是一种存储k维空间中数据的平衡二叉树型结构，主要用于范围搜索和最近邻搜索。\
  kd树实质是一种空间划分树，其每个节点对应一个k维的点，\
  每个非叶节点相当于一个分割超平面，将其所在区域划分为两个子区域。\
"""


class KDNode(object):
    def __init__(
            self, features, label,
            left=None, right=None,
            axis_no=0, depth=0
    ):
        # 每个节点的特征值
        self.features = features
        # 每个节点对应的标签
        self.label = label
        # 节点的左孩子
        self.left = left
        # 节点的右孩子
        self.right = right
        # 划分维度编号
        self.axis_no = axis_no
        # 节点所在的深度
        #self.depth = depth

"""
构造KD树
1、构造根节点：选取划分轴内所有点的中位数
2、计算深度为j的节点所对应的划分轴
3、重复上述步骤，直到无其他实例点
"""

class KDTree():
    def __init__(self,n_features):
        #根节点
        self.root = None
        # 特征个数
        self.dimensions = n_features
        # 最近保存的k个节点
        self.k_neighbour = []
        # 记录访问过的几点的路径
        self.path = []
    
    # 计算划分轴
    # 对于n_features个特征的数据集，深度为j的节点，\
    # 选择L为划分轴
    # L = j mod k

    def getAxis(self,depth,n_features):
        return depth % n_features

    """
     Examples
    -------
    a = np.array([
        [300, 1, 10, 999],
        [200, 2, 30, 987],
        [100, 3, 10, 173]
    ])
    sort_index = np.argsort(a, axis=0)
    print(sort_index)
    >>>
        [[2 0 0 2]
         [1 1 2 1]
         [0 2 1 0]]
         
    sort_index[:, 3]
    >>> 
        [2 1 0]
    print(a[sort_index[:, 3]])
    >>>
        [[100   3  10 173]
         [200   2  30 987]
         [300   1  10 999]]
    """
    # 根据划分轴计算中位数
    def getArgSortMatrix(self,dataset,label,depth):
        sort_index = np.argsort(dataset, axis=0)
        return dataset[sort_index[:,depth]],label[sort_index[:,depth]]


    # 欧式距离
    def getDistance(self,node,target):
        diff = (node.features - target)**2
        sqrt_distance = diff.sum(axis=1)
        return sqrt_distance ** 0.5

    """
        1. 如果数据集中只有一条数据，则赋予空的叶子节点
        2. 如果不止一条数据，则进行如下操作：
            a. 根据构造树当前的深度，选定划分轴（根据哪个特征进行划分）
            b. 根据划分轴（该特征），对数据集按照该特征从小到大排序
            c. 选出中位数、排序特征中大于、小于中位数的子数据集
            d. 递归调用自身，构造KDTree
    """
    def createKDTree(self,dataset,label,depth):
        samples = dataset.shape[0]
        if samples < 1:
            return None
        if samples == 1:
            new_node = KDNode(dataset[0],label[0],depth=depth)
            return  new_node
        else:
            dataset,label = self.getArgSortMatrix(dataset,label,depth)
            middle = int(samples / 2)
            new_node = KDNode(dataset[middle],label[middle],depth=depth)

            left_dataset = dataset[:middle]
            left_label = label[:middle]
            new_node.left = self.createKDTree(left_dataset,left_label,depth+1)
            
            if(middle+1 < samples):
                right_dataset = dataset[middle+1:]
                right_label = label[middle+1:]
                new_node.right = self.createKDTree(right_dataset,right_label,depth+1)
        
        return new_node

    def getKNode(self, kdTree, features,k):
        if kdTree is None:
            return
        else:
            # 计算一下距离，在下一步中看看是不是真的要走这个分支
            distance = self.getDistance(kdTree, features)
            # 因为维护的是最小堆，每个新节点的距离需要和最小堆中的最大值进行比较，所以放在前面了
            self.k_neighbour.reverse()
            if (len(self.k_neighbour) < k) or (distance < self.k_neighbour[0]['distance']):
                # 先顺着左边走到底，然后在看右边的
                self.getKNode(kdTree.left, features, k)
                self.getKNode(kdTree.right, features, k)
                # 把走过的节点都记录下来
                self.path.append(kdTree)
                # 3.将该点加入堆中，维护一个最小堆，按照distance进行排序
                self.k_neighbour.append({
                    'node': kdTree,
                    'distance': distance
                })
                self.k_neighbour = heapq.nsmallest(
                    k, self.k_neighbour, key=lambda s: s['distance']
                )
                #self.k_neighbour = self.getKheadNode(k, self.k_neighbour, key=lambda s: s['distance'])
            return


"""
KD树搜索
0、设 L 为一个有 k 个空位的列表，用于保存已搜寻到的最近点。

1、根据 p 的坐标值和每个节点的切分向下搜索（也就是说，如果树的节点是按照 xr=a 进行切分，并且 p 的 r 坐标小于 a，则向左枝进行搜索；反之则走右枝）。

2、当达到一个底部节点时，将其标记为访问过。如果 L 里不足 k 个点，则将当前节点的特征坐标加入 L ；如果 L 不为空并且当前节点的特征与 p 的距离小于 L 里最长的距离，则用当前特征替换掉 L 中离 p 最远的点。

3、如果当前节点不是整棵树最顶端节点，执行 (a)；反之，输出 L，算法完成。

    a. 向上爬一个节点。如果当前（向上爬之后的）节点未曾被访问过，将其标记为被访问过，然后执行 (1) 和 (2)；如果当前节点被访问过，再次执行 (a)。

        1. 如果此时 L 里不足 k 个点，则将节点特征加入 L；如果 L 中已满 k 个点，且当前节点与 p 的距离小于 L 里最长的距离，则用节点特征替换掉 L 中离最远的点。

        2. 计算 p 和当前节点切分线的距离。如果该距离大于等于 L 中距离 p 最远的距离，则在切分线另一边不会有更近的点，执行(三)；如果该距离小于 L 中最远的距离，则切分线另一边可能有更近的点，因此在当前节点的另一个枝从 (一) 开始执行。
"""




if __name__ == '__main__':
    dataset = np.array([
        [2,3],
        [5,4],
        [9,6],
        [4,7],
        [8,1],
        [7,2]
    ])

    labels = np.array([1,1,-1,1,1,-1])

    tree = KDTree(2)

    nodeTree = tree.createKDTree(dataset,labels,0)

    node = np.array([[10,12]])

    tree.getKNode(nodeTree,node,2)
    print(tree.k_neighbour)
    